# -*- coding: utf-8 -*-

# 导入pyspark
from pyspark import SparkContext 
from pyspark.mllib.regression import LabeledPoint
import numpy as np

# 导入线性回归
from pyspark.mllib.regression import LinearRegressionWithSGD
from pyspark.mllib.tree import DecisionTree

# 导入可视化
import matplotlib
matplotlib.use('Agg') # 不回显
import matplotlib.pyplot as plt

# 函数-决策树模型创建特征向量
def extract_features_dt(record):
	return np.array(map(float, record[1:4]))

# 函数-提取实际租借车数
def extract_label(record):
	return float(record[4])

# 函数-平方对数误差
def squared_log_error(pred, actual):
	return (np.log(pred + 1) - np.log(actual + 1))**2


# 函数-参数设置对决策树性能的影响
def evaluate_dt(train, test, maxDepth, maxBins):
	
	# 训练DecisionModel
	model = DecisionTree.trainRegressor(train, {}, impurity = 'variance', maxDepth = maxDepth, maxBins = maxBins)
	preds = model.predict(test.map(lambda p: p.features))
	actual = test.map(lambda p: p.label)

	# 返回rsmle
	tp = actual.zip(preds)
	rmsle = np.sqrt(tp.map(lambda (t, p): squared_log_error(t,p)).mean())
	return rmsle


# 函数-画出比较
def draw_figure(actual, predict, title):

	figName = "%s.png"%title
	# 基本图构成
	plt.figure(num=1, figsize=(8,6))
	plt.title(title, size=14)
	plt.xlabel('x', size=14)
	plt.ylabel('y', size=14)

	# 点图
	plt.plot(actual, predict)

	# 保存图片
	plt.savefig(figName, format='png')


# main函数部分
sc = SparkContext("yarn-client", "Decision Tree Regression Quality Spark App")

# 加载数据集
path = "hdfs://192.168.119.141:9100/data/materialConductive/regressionConductive.csv" # 数据集位置
raw_data = sc.textFile(path) # 数据集
records = raw_data.map(lambda x: x.split(",")) # 通过逗号分隔
records.cache() # 缓存rdd

# 决策树模型
data_dt = records.map(lambda r: LabeledPoint(extract_label(r), extract_features_dt(r)))

# train/test 数据准备
data_with_idx_dt = data_dt.zipWithIndex().map(lambda (k, v): (v, k))
test_dt = data_with_idx_dt.sample(False, 0.2, 42)
train_dt = data_with_idx_dt.subtractByKey(test_dt)
train_data_dt = train_dt.map(lambda (idx, p): p)
test_data_dt = test_dt.map(lambda (idx, p): p)

depthParams = [1, 2, 3, 4, 5, 10, 20]
binsParams = [2, 4, 8, 16, 32, 64]
depthMetrics = [evaluate_dt(train_data_dt, test_data_dt, maxDepth, 16) for maxDepth in depthParams]
binsMetrics = [evaluate_dt(train_data_dt, test_data_dt, 5, maxBins) for maxBins in binsParams]

draw_figure(depthParams, depthMetrics, "DecisionTreeConductiveQualityDepth")
draw_figure(binsParams, binsMetrics, "DecisionTreeConductiveQualityBins")

# 关闭sc
sc.stop()
